﻿using System;
using System.Collections.Generic;
using System.Linq;
using MethodCaching.Weaving.Extensions;
using Mono.Cecil;
using Mono.Cecil.Cil;
using Mono.Cecil.Rocks;
namespace MethodCaching.Weaving
{
    /// <summary>
    /// Implements a <see cref="MethodCache"/> property 
    /// that can be used to invoke cache decorated methods.
    /// </summary>
    public class TypeWeaver : ITypeWeaver
    {
        private const string STATIC_METHOD_CACHE_FIELD_NAME = "__staticMethodCache";
        private const string INSTANCE_METHOD_CACHE_FIELD_NAME = "__methodCache";

        private const string STATIC_METHOD_CACHE_PROPERTY_NAME = "__StaticMethodCache";
        private const string INSTANCE_METHOD_CACHE_PROPERTY_NAME = "__MethodCache";

        private const string STATIC_METHOD_CACHE_GETTER_METHOD_NAME = "get___StaticMethodCache";
        private const string INSTANCE_METHOD_CACHE_GETTER_METHOD_NAME = "get___MethodCache";

        private readonly IMethodSelector _methodSelector;                        
        private IEnumerable<MethodDefinition> _methods;
        private TypeDefinition _typeDefinition;
        private ModuleDefinition _moduleDefinition;
        private TypeReference _methodCacheTypeReference;
        private MethodReference _methodCacheConstructorReference;
        
        private FieldDefinition _instanceMethodCacheFieldDefinition;
        private FieldDefinition _staticMethodCacheFieldDefinition;
                                
        private MethodDefinition _instanceMethodCacheGetterMethodDefinition;        
        private MethodDefinition _staticMethodCacheGetterMethodDefinition;

        /// <summary>
        /// Initializes a new instance of the <see cref="TypeWeaver"/> class.
        /// </summary>
        /// <param name="methodSelector">The <see cref="IMethodSelector"/> that is responsible for </param>
        public TypeWeaver(IMethodSelector methodSelector)        
        {
            _methodSelector = methodSelector;
        }

        /// <summary>
        /// Weaves the target type
        /// </summary>
        /// <param name="typeDefinition">The target type.</param>
        public void Weave(TypeDefinition typeDefinition)
        {
            Initialize(typeDefinition);                        
            ResolveTargetMethods();            
            WeaveMethods();            
        }

        private void Initialize(TypeDefinition typeDefinition)
        {
            _instanceMethodCacheFieldDefinition = null;
            _staticMethodCacheFieldDefinition = null;
            _instanceMethodCacheGetterMethodDefinition = null;
            _staticMethodCacheGetterMethodDefinition = null;
            _methods = null;
            _typeDefinition = typeDefinition;
            _moduleDefinition = _typeDefinition.Module;
            _methodCacheConstructorReference = _moduleDefinition.Import(typeof(MethodCache).GetConstructor(Type.EmptyTypes));
        }


        private void ImplementInstanceMethodCache()
        {
            if (_instanceMethodCacheGetterMethodDefinition != null)
                return;
            PropertyDefinition methodCacheProperty = CreateInstanceMethodCachePropertyWithBackingField();            
                methodCacheProperty.GetMethod = EmitInstanceMethodCacheGetterMethod();            
        }

        private void ImplementStaticMethodCache()
        {
            if (_staticMethodCacheGetterMethodDefinition != null)
                return;
            PropertyDefinition methodCacheProperty = CreateStaticMethodCachePropertyWithBackingField();            
            methodCacheProperty.GetMethod = EmitStaticMethodCacheGetterMethod();            
        }
               

        private PropertyDefinition CreateInstanceMethodCachePropertyWithBackingField()
        {
            CreateInstanceMethodCacheBackingField();
            return CreateInstanceMethodCachePropertyDefinition();
        }

        private PropertyDefinition CreateStaticMethodCachePropertyWithBackingField()
        {
            CreateStaticMethodCacheBackingField();
            return CreateStaticMethodCachePropertyDefinition();
        }



        private void CreateInstanceMethodCacheBackingField()
        {
            _methodCacheTypeReference = _moduleDefinition.Import(typeof(MethodCache));
            const FieldAttributes fieldAttributes = FieldAttributes.Private;
            
            _instanceMethodCacheFieldDefinition = new FieldDefinition(INSTANCE_METHOD_CACHE_FIELD_NAME, fieldAttributes,
                                                                  _methodCacheTypeReference);
            _typeDefinition.Fields.Add(_instanceMethodCacheFieldDefinition);
            
        }

        private void CreateStaticMethodCacheBackingField()
        {
            _methodCacheTypeReference = _moduleDefinition.Import(typeof(MethodCache));
            const FieldAttributes fieldAttributes = FieldAttributes.Private | FieldAttributes.Static;
            _staticMethodCacheFieldDefinition = new FieldDefinition(STATIC_METHOD_CACHE_FIELD_NAME, fieldAttributes,
                                                                  _methodCacheTypeReference);
            _typeDefinition.Fields.Add(_staticMethodCacheFieldDefinition);

        }


        

        private PropertyDefinition CreateInstanceMethodCachePropertyDefinition()
        {            
            var methodCachePropertyDefinition = new PropertyDefinition(INSTANCE_METHOD_CACHE_PROPERTY_NAME, PropertyAttributes.None,
                                                                    _methodCacheTypeReference);            
            _typeDefinition.Properties.Add(methodCachePropertyDefinition);
            return methodCachePropertyDefinition;
        }

        private PropertyDefinition CreateStaticMethodCachePropertyDefinition()
        {
            var methodCachePropertyDefinition = new PropertyDefinition(STATIC_METHOD_CACHE_PROPERTY_NAME, PropertyAttributes.None,
                                                                    _methodCacheTypeReference) {HasThis = true};
            methodCachePropertyDefinition.HasThis = false;
            _typeDefinition.Properties.Add(methodCachePropertyDefinition);
            return methodCachePropertyDefinition;
        }

        

        private MethodDefinition EmitInstanceMethodCacheGetterMethod()
        {
            const MethodAttributes methodAttributes = MethodAttributes.Private | MethodAttributes.SpecialName |
                                                      MethodAttributes.HideBySig;
                                    
            var methodDefinition = new MethodDefinition(INSTANCE_METHOD_CACHE_GETTER_METHOD_NAME,
                                                        methodAttributes, _methodCacheTypeReference)
                                       {Body = {InitLocals = true}};

            methodDefinition.Body.Variables.Add(new VariableDefinition(_methodCacheTypeReference));
            methodDefinition.Body.Variables.Add(new VariableDefinition(_moduleDefinition.Import(typeof(bool))));
            ILProcessor ilProcessor =  methodDefinition.Body.GetILProcessor();                        
            Instruction firstLabel = ilProcessor.Create(OpCodes.Ldarg_0);            
            Instruction secondLabel = ilProcessor.Create(OpCodes.Ldloc_0);
            ilProcessor.Emit(OpCodes.Nop);
            ilProcessor.Emit(OpCodes.Ldarg_0);
            ilProcessor.Emit(OpCodes.Ldfld, _instanceMethodCacheFieldDefinition);
            ilProcessor.Emit(OpCodes.Ldnull);
            ilProcessor.Emit(OpCodes.Ceq);
            ilProcessor.Emit(OpCodes.Ldc_I4_0);
            ilProcessor.Emit(OpCodes.Ceq);
            ilProcessor.Emit(OpCodes.Stloc_1);
            ilProcessor.Emit(OpCodes.Ldloc_1);
            ilProcessor.Emit(OpCodes.Brtrue_S, firstLabel);
            ilProcessor.Emit(OpCodes.Ldarg_0);
            ilProcessor.Emit(OpCodes.Newobj, _methodCacheConstructorReference);
            ilProcessor.Emit(OpCodes.Stfld, _instanceMethodCacheFieldDefinition);
            ilProcessor.Append(firstLabel);
            ilProcessor.Emit(OpCodes.Ldfld, _instanceMethodCacheFieldDefinition);
            ilProcessor.Emit(OpCodes.Stloc_0);
            ilProcessor.Emit(OpCodes.Br_S, secondLabel);
            ilProcessor.Append(secondLabel);
            
            ilProcessor.Emit(OpCodes.Ret);
            _typeDefinition.Methods.Add(methodDefinition);
            _instanceMethodCacheGetterMethodDefinition =  methodDefinition;
            return methodDefinition;
        }

        private MethodDefinition EmitStaticMethodCacheGetterMethod()
        {
            const MethodAttributes methodAttributes = MethodAttributes.Private | MethodAttributes.SpecialName |
                                                      MethodAttributes.HideBySig | MethodAttributes.Static;

            var methodDefinition = new MethodDefinition(STATIC_METHOD_CACHE_GETTER_METHOD_NAME,
                                                        methodAttributes, _methodCacheTypeReference)
                                       {Body = {InitLocals = true}};

            ILProcessor ilProcessor = methodDefinition.Body.GetILProcessor();
            Instruction endLabel = ilProcessor.Create(OpCodes.Nop);                        
            //Load the field onto the stack
            ilProcessor.Emit(OpCodes.Ldsfld, _staticMethodCacheFieldDefinition);
            //Load the value of null of the stack
            ilProcessor.Emit(OpCodes.Ldnull);
            //Compare the field reference with null
            //If it is NULL, the value of 1 is pushed onto the stack, otherwise 0
            ilProcessor.Emit(OpCodes.Ceq);            
            //If the value currently on stack is 0 (false), branch to the first label.
            ilProcessor.Emit(OpCodes.Brfalse_S, endLabel);
            //else create a new MethodCache instance
            ilProcessor.Emit(OpCodes.Newobj, _methodCacheConstructorReference);
            //Store the new instance in the field reference
            ilProcessor.Emit(OpCodes.Stsfld, _staticMethodCacheFieldDefinition);
            ilProcessor.Append(endLabel);
            //Load the field referenc onto the stack
            ilProcessor.Emit(OpCodes.Ldsfld, _staticMethodCacheFieldDefinition);
            //Return
            ilProcessor.Emit(OpCodes.Ret);
            _typeDefinition.Methods.Add(methodDefinition);
            _staticMethodCacheGetterMethodDefinition = methodDefinition;
            return methodDefinition;
        }



        private void WeaveMethods()
        {
            foreach (var methodDefinition in _methods)
                WeaveMethod(methodDefinition);
        }

        private void WeaveMethod(MethodDefinition methodDefinition)
        {
            if (methodDefinition.IsStatic)
            {
                WeaveStaticMethod(methodDefinition);
            }
            else
                WeaveInstanceMethod(methodDefinition);
        }


        private void WeaveInstanceMethod(MethodDefinition methodDefinition)
        {            
            ImplementInstanceMethodCache();
            var extractedMethod = ExtractMethod(methodDefinition);
            ClearMethodBody(methodDefinition);
            EmitNewInstanceMethodBody(methodDefinition, extractedMethod);
        }

        private void WeaveStaticMethod(MethodDefinition methodDefinition)
        {           
            ImplementStaticMethodCache();
            var extractedMethod = ExtractMethod(methodDefinition);
            ClearMethodBody(methodDefinition);
            EmitNewStaticMethodBody(methodDefinition, extractedMethod);
        }


        

        private void EmitNewInstanceMethodBody(MethodDefinition methodDefinition, MethodDefinition extractedMethod)
        {
            TypeReference[] genericArgs = GetGenericArguments(methodDefinition);
                                    
            ILProcessor ilProcessor =  methodDefinition.Body.GetILProcessor();                                            
            ilProcessor.Emit(OpCodes.Nop);
            ilProcessor.Emit(OpCodes.Ldarg_0);
            ilProcessor.Emit(OpCodes.Call,_instanceMethodCacheGetterMethodDefinition);
            ilProcessor.Emit(OpCodes.Ldarg_0);
            ilProcessor.Emit(OpCodes.Ldftn, extractedMethod);            
            ilProcessor.Emit(OpCodes.Newobj,CreateClosedGenericFuncConstructor(genericArgs));
            for (int index = 0; index < genericArgs.Length -1; index++)
                ilProcessor.Emit(OpCodes.Ldarg, index + 1);
            ilProcessor.Emit(OpCodes.Callvirt, CreateInvokeMethod(genericArgs));                        
            ilProcessor.Emit(OpCodes.Ret);
        }

        private void EmitNewStaticMethodBody(MethodDefinition methodDefinition, MethodDefinition extractedMethod)
        {
            TypeReference[] genericArgs = GetGenericArguments(methodDefinition);

            ILProcessor ilProcessor = methodDefinition.Body.GetILProcessor();
            ilProcessor.Emit(OpCodes.Nop);              
            ilProcessor.Emit(OpCodes.Call, _staticMethodCacheGetterMethodDefinition);
            ilProcessor.Emit(OpCodes.Ldnull);
            ilProcessor.Emit(OpCodes.Ldftn, extractedMethod);
            ilProcessor.Emit(OpCodes.Newobj, CreateClosedGenericFuncConstructor(genericArgs));
            for (int index = 0; index < genericArgs.Length - 1; index++)
                ilProcessor.Emit(OpCodes.Ldarg, index);
            ilProcessor.Emit(OpCodes.Callvirt, CreateInvokeMethod(genericArgs));
            ilProcessor.Emit(OpCodes.Ret);
        }


        private static void ClearMethodBody(MethodDefinition methodDefinition)
        {
            methodDefinition.Body.Instructions.Clear();
            methodDefinition.Body.Variables.Clear();
        }

        private TypeReference GetFuncTypeReference(int parameterCount)
        {            
            switch (parameterCount)
            {
                case 0:
                    return _moduleDefinition.Import(typeof(Func<>));
                case 1:
                    return _moduleDefinition.Import(typeof(Func<,>));
                case 2:
                    return _moduleDefinition.Import(typeof(Func<,,>));
                case 3:
                    return _moduleDefinition.Import(typeof(Func<,,,>));
                case 4:
                    return _moduleDefinition.Import(typeof(Func<,,,,>));
                default:
                    throw new ArgumentOutOfRangeException("parameterCount","Cannot create function delegates for more than 4 parameters");
            }
        }



        private MethodReference CreateClosedGenericFuncConstructor(TypeReference[] genericArgs)
        {                                   

            TypeReference funcReference = GetFuncTypeReference(genericArgs.Length - 1);

            GenericInstanceType genericInstanceType = funcReference.MakeGenericInstanceType(genericArgs);
            
            TypeDefinition resolvedGenericType =  genericInstanceType.Resolve();

            MethodDefinition openGenericConstructor = resolvedGenericType.GetConstructors().First();

            MethodReference gctor = _moduleDefinition.Import(openGenericConstructor.MakeHostInstanceGeneric(genericArgs));

            return gctor;
        }


        private static TypeReference[] GetGenericArguments(MethodDefinition methodDefinition)
        {
            var genericArgs = new TypeReference[methodDefinition.Parameters.Count + 1];
            for (int index = 0; index < methodDefinition.Parameters.Count; index++)
            {
                var parameterDefinition = methodDefinition.Parameters[index];
                genericArgs[index] = parameterDefinition.ParameterType;
            }
            genericArgs[genericArgs.Length -1] = methodDefinition.ReturnType;
            return genericArgs;
        }


        private MethodReference CreateInvokeMethod(TypeReference[] genericArgs)
        {
            MethodReference invokeMethodReference = GetInvokeMethodReference(genericArgs);            
            return _moduleDefinition.Import(invokeMethodReference.MakeGenericInstanceMethod(genericArgs));
        }

        private MethodReference GetInvokeMethodReference(TypeReference[] genericArgs)
        {
            var method =
                typeof (MethodCache).GetMethods().Where(
                    m => m.Name == "Invoke" && m.GetParameters().Count() == genericArgs.Length).
                    FirstOrDefault();

            return _moduleDefinition.Import(method);
        }


        private MethodDefinition ExtractMethod(MethodDefinition methodDefinition)
        {
            MethodDefinition newMethod;
            if (methodDefinition.IsStatic)
                newMethod = CreateExtractedStaticMethod(methodDefinition);
            else
                newMethod = CreateExtractedInstanceMethod(methodDefinition);
            return CopyMethodBody(methodDefinition, newMethod);
        }
                
        private MethodDefinition CreateExtractedInstanceMethod(MethodDefinition methodDefinition)
        {
            const MethodAttributes methodAttributes = MethodAttributes.HideBySig | MethodAttributes.Private;
            return new MethodDefinition(GetMethodName(methodDefinition), methodAttributes, _typeDefinition.Module.Import(methodDefinition.ReturnType));
        }

        private MethodDefinition CreateExtractedStaticMethod(MethodDefinition methodDefinition)
        {
            const MethodAttributes methodAttributes = MethodAttributes.HideBySig | MethodAttributes.Private | MethodAttributes.Static;
            return new MethodDefinition(GetMethodName(methodDefinition), methodAttributes, _typeDefinition.Module.Import(methodDefinition.ReturnType));
        }

        private MethodDefinition CopyMethodBody(MethodDefinition methodDefinition, MethodDefinition newMethod)
        {
            newMethod.Body.InitLocals = methodDefinition.Body.InitLocals;

            foreach (var variableDefinition in methodDefinition.Body.Variables)
                newMethod.Body.Variables.Add(
                    new VariableDefinition(_typeDefinition.Module.Import(variableDefinition.VariableType)));

            foreach (var parameterDefinition in methodDefinition.Parameters)
                newMethod.Parameters.Add(
                    new ParameterDefinition(_typeDefinition.Module.Import(parameterDefinition.ParameterType)));

            foreach (var instruction in methodDefinition.Body.Instructions)
                newMethod.Body.Instructions.Add(instruction);

            _typeDefinition.Methods.Add(newMethod);
            return newMethod;
        }

        private static string GetMethodName(MethodDefinition methodDefinition)
        {
            return string.Format("__Do{0}", methodDefinition.Name);
        }
      
        private void ResolveTargetMethods()
        {
            _methods = _methodSelector.Select(_typeDefinition);                        
        }        
    }
}
